/*
*         (C) COPYRIGHT NBI Limited
*              ALL RIGHT RESERVED
*
* File        : fastflow/protobuf/onnx_proto_helper.cc
* Authors     : dzhang
* Create Time : 2021/09/03 09:28:23
* Description :
*
*/

#include "fastflow/graph/onnx_proto_helper.h"
#include <stdio.h>
#include <stdint.h>
#include <fstream>
#include "onnx/shape_inference/implementation.h"

namespace fastflow{

bool onnx_read_proto_from_binary(const char* filepath, google::protobuf::Message* message) {
    std::ifstream fs(filepath, std::ifstream::in | std::ifstream::binary);
    if (!fs.is_open()) {
        fprintf(stderr, "open failed %s\n", filepath);
        return false;
    }   

    google::protobuf::io::IstreamInputStream input(&fs);
    google::protobuf::io::CodedInputStream codedstr(&input);
#if GOOGLE_PROTOBUF_VERSION >= 3011000
    codedstr.SetTotalBytesLimit(INT_MAX);
#else
    codedstr.SetTotalBytesLimit(INT_MAX, INT_MAX/2);
#endif

    bool success = message->ParseFromCodedStream(&codedstr);

    fs.close();

    return success;
}
bool onnx_write_proto_from_binary(const char* filepath, const google::protobuf::Message* message) {
    std::ofstream fs(filepath);
    if (fs.fail()) {
        fprintf(stderr, "open failed %s\n", filepath);
        return false;
    }   
    message->SerializeToOstream(&fs);
    fs.close();
    return true;
}

bool onnx_shape_inference(ONNX_NAMESPACE::ModelProto& m){
    ONNX_NAMESPACE::shape_inference::InferShapes(m);
    return true;
}

}  // namespace fastflow